package dd.lo.pdfimg;

import dd.lo.HelloCV;
import dd.lo.util.ImageUtils;
import javafx.application.Application;
import javafx.application.Platform;
import javafx.beans.property.SimpleStringProperty;
import javafx.concurrent.Task;
import javafx.scene.Scene;
import javafx.scene.control.Button;
import javafx.scene.control.TextField;
import javafx.scene.image.Image;
import javafx.scene.image.ImageView;
import javafx.scene.layout.BorderPane;
import javafx.scene.layout.FlowPane;
import javafx.scene.layout.VBox;
import javafx.scene.text.Text;
import javafx.stage.FileChooser;
import javafx.stage.Modality;
import javafx.stage.Stage;
import org.apache.commons.lang3.math.NumberUtils;
import org.apache.pdfbox.cos.COSName;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDPage;
import org.apache.pdfbox.pdmodel.PDResources;
import org.apache.pdfbox.pdmodel.common.PDStream;
import org.apache.pdfbox.pdmodel.graphics.PDXObject;
import org.apache.pdfbox.pdmodel.graphics.image.PDImageXObject;
import org.opencv.core.Core;
import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.core.MatOfByte;
import org.opencv.core.MatOfPoint;
import org.opencv.core.Point;
import org.opencv.core.Rect;
import org.opencv.core.Scalar;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

public class PdfImgSplitterDisplay extends Application {

    public static void main(String[] args) {
        System.load(HelloCV.class.getResource("/libopencv_java481.dylib").getPath());
        PdfImgSplitterDisplay.start(args);
    }

    public static void start(String[] args) {
        launch(args);
    }

    //轮廓与该外接矩形的面积比，如果大于阈值，则认为轮廓已经很接近矩形，该轮廓可以被认为用于切图
    private static final double EXTENT_THRESHOLD = 0.95;

    private static final int DEFAULT_WIDTH = 1080;
    private static final int DEFAULT_HEIGHT = 768;

    private static final int SHOW_IMG_WIDTH = 800;
    private static final int SHOW_IMG_HEIGHT = 600;

    private double offSet = 50;

    private File file;

    private FlowPane imageFlowPane;

    private SimpleStringProperty pageNo;

    private SimpleStringProperty imageNo;

    //根据聂鹏的实验结果。对于一个区域，如果(短边/区域面积) < MIN_REGION_AREA_RATIO，这个区域需要被抹掉
    private static final double MIN_REGION_AREA_RATIO = 30;

    private void showImageNewWindow(Mat img, String winTitle) {
        Mat dst = new Mat();
        img.copyTo(dst);
        Platform.runLater(() -> {
            ImageView view = new ImageView();
            //缩放、转码
            ImageUtils.fitImgSize(dst, SHOW_IMG_WIDTH, SHOW_IMG_HEIGHT);
            MatOfByte buffer = new MatOfByte();
            Imgcodecs.imencode(".jpeg", dst, buffer);
            view.setImage(new Image(new ByteArrayInputStream(buffer.toArray())));
            //贴到组件上
            VBox vBox = new VBox();
            vBox.getChildren().addAll(view);
            Scene scene = new Scene(vBox, SHOW_IMG_WIDTH, SHOW_IMG_HEIGHT);
            Stage newWindow = new Stage();
//        newWindow.initModality(Modality.APPLICATION_MODAL); // 设置此模式，则窗口打开时，其他窗口无法操作
            newWindow.initModality(Modality.NONE);
            newWindow.setTitle(winTitle);
            newWindow.setScene(scene);
            newWindow.setX(offSet);
            newWindow.setY(offSet);
            offSet += 50;
            newWindow.show();
        });
    }

    private List<MatOfPoint> findOuterContours(Mat image) {
        List<MatOfPoint> contours = new ArrayList<>();
        Mat hierarchy = new Mat();
        Imgproc.findContours(image, contours, hierarchy, Imgproc.RETR_TREE, Imgproc.CHAIN_APPROX_SIMPLE);
        List<MatOfPoint> outerContours = new ArrayList<>(contours.size());
        for (int i = 0; i < contours.size(); ++i) {
            MatOfPoint contour = contours.get(i);
            //0~3分别表示后一个轮廓、前一个轮廓、父轮廓、内嵌轮廓的索引值
            //参考资料：https://blog.csdn.net/yuanren201/article/details/104418710
            double[] h = hierarchy.get(0, i);
            if (h[3] == 0.0) {
                //内嵌轮廓的索引值为0时，说明这个轮廓是最外围的
                outerContours.add(contour);
            }
        }
        return outerContours;
    }

    private Mat splitAndDraw(Mat src) {
        Mat grayImg = new Mat();
        Imgproc.cvtColor(src, grayImg, Imgproc.COLOR_BGR2GRAY);
        Core.normalize(grayImg, grayImg, 0, 255, Core.NORM_MINMAX);
        //预处理获得原始切图Mask
        Mat splitMask = new Mat();
        Imgproc.threshold(grayImg, splitMask, 220, 255, Imgproc.THRESH_BINARY);
        showImageNewWindow(splitMask, "预处理的Mask");
        //给图像加上外围的白边，让涂黑的轮廓不会粘到图片边缘导致不能被识别
        Core.copyMakeBorder(splitMask, splitMask, 10, 10, 10, 10, Core.BORDER_CONSTANT, new Scalar(255));
        showImageNewWindow(splitMask, "加上外围边框后的Mask");
        //第一次找出所有区域
        List<MatOfPoint> contours = findOuterContours(splitMask);
        //抹掉过小的区域
        for (MatOfPoint contour : contours) {
            Rect region = Imgproc.boundingRect(contour);
            double a = Math.min(src.width(), src.height()) / MIN_REGION_AREA_RATIO;
            double minArea = a * a;
            if (region.area() < minArea) {
                MatOfPoint points = new MatOfPoint(
                        region.tl(),
                        new Point(region.x + region.width, region.y),
                        region.br(),
                        new Point(region.x, region.y + region.height)
                );
                Imgproc.fillPoly(splitMask, Collections.singletonList(points), new Scalar(255, 255, 255));
            }
        }
        showImageNewWindow(splitMask, "抹掉过小的区域后的Mask");
        //至多重复做一定次数非矩形区域填充，让其轮廓尽量与其外接矩形一致
        int turn = 0;
        double regionExtend = 0;
        while (turn < 5 && regionExtend < 0.95) {
            //填充非矩形区域
            double minExtent = EXTENT_THRESHOLD;
            contours = findOuterContours(splitMask);
            for (MatOfPoint contour : contours) {
                Rect rect = Imgproc.boundingRect(contour);
                double contourArea = Imgproc.contourArea(contour);
                double extent = contourArea / rect.area();
                minExtent = Math.min(minExtent, extent);
                if (extent > EXTENT_THRESHOLD) {
                    continue;
                }
                //填充外接矩形
                MatOfPoint points = new MatOfPoint(
                        rect.tl(),
                        new Point(rect.x + rect.width, rect.y),
                        rect.br(),
                        new Point(rect.x, rect.y + rect.height)
                );
                Imgproc.fillPoly(splitMask, Collections.singletonList(points), new Scalar(0, 0, 0));
            }
            regionExtend = minExtent;
            System.out.printf("当前最小extend值%f\n", minExtent);
            showImageNewWindow(splitMask, String.format("%d次填充非矩形区域", turn + 1));
            turn++;
        }
        //设定展示的图像
        contours = findOuterContours(splitMask);
        Mat outputImg = new Mat();
        splitMask.copyTo(outputImg);
        ImageUtils.adaptImgChannel(outputImg);
        for (int i = 0; i < contours.size(); ++i) {
            MatOfPoint contour = contours.get(i);
            Rect rect = Imgproc.boundingRect(contour);
            //轮廓用绿色线条展示出来
            Imgproc.polylines(outputImg, contours, true, new Scalar(0, 255, 0), 1, Imgproc.LINE_4);
            //外接矩形ROI用蓝色线条展示出来
            Imgproc.rectangle(outputImg, rect, new Scalar(255, 0, 0), 4);
        }
        return outputImg;
    }

    private void pdfProcess(PDDocument document) {
        if (null == document) {
            System.out.println("文件读取失败");
            return;
        }
        int pNo = NumberUtils.toInt(pageNo.getValue(), 0);
        int iNo = NumberUtils.toInt(imageNo.getValue(), 0);
        Task<Void> task = new Task<Void>() {
            @Override
            protected Void call() {
                int imageNumber;
                try {
                    PDPage page = document.getPage(pNo);
                    if (null == page) {
                        System.out.printf("没有第%d页\n", pNo);
                        return null;
                    }
                    PDResources resources = page.getResources();
                    imageNumber = 0;
                    for (COSName cosName : resources.getXObjectNames()) {
                        PDXObject xObject;
                        try {
                            xObject = resources.getXObject(cosName);
                        } catch (IOException e) {
                            continue;
                        }
                        if (xObject instanceof PDImageXObject) {
                            if (imageNumber < iNo) continue;
                            PDStream stream = xObject.getStream();
                            try {
                                //先转一次流才能拿到真正的图片对象
                                PDImageXObject imageObject = new PDImageXObject(stream, resources);
                                //获得awt的BufferedImage
                                BufferedImage image = imageObject.getImage();
                                //将BufferedImage利用ImageIO.write按jpeg格式输出到ByteArrayOutputStream流中
                                ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
                                ImageIO.write(image, "jpeg", byteArrayOutputStream);
                                byteArrayOutputStream.flush();
                                //此时就可以得到图片的byte信息，在ByteArrayOutputStream中，然后就可以通过byte信息将图片load到Mat中
                                Mat img = Imgcodecs.imdecode(new MatOfByte(byteArrayOutputStream.toByteArray()), Imgcodecs.IMREAD_UNCHANGED);
                                if (img.empty()) continue;
                                if (img.width() < 50 || img.height() < 50) continue; //太小的图片忽略掉
                                img.convertTo(img, CvType.CV_8UC3);
                                if (img.channels() < 3) {
                                    //如果图片不是3通道的，先自己合并自己变成3通道，否则后续步骤处理报错
                                    Core.merge(Arrays.asList(img, img, img), img);
                                }
                                //寻找图片中的小图，展示切割图片的轮廓外矩形
                                Mat outputImg = splitAndDraw(img);
                                ImageUtils.fitImgSize(outputImg, DEFAULT_WIDTH, DEFAULT_HEIGHT);
                                MatOfByte matOfByte = new MatOfByte();
                                Imgcodecs.imencode(".jpeg", outputImg, matOfByte);
                                ImageView imageView = new ImageView();
                                imageView.setImage(new Image(new ByteArrayInputStream(matOfByte.toArray())));
                                Platform.runLater(() -> {
                                    imageFlowPane.getChildren().clear();
                                    imageFlowPane.getChildren().add(imageView);
                                });
                                break;
                            } catch (IOException ignored) {
                            }
                            imageNumber++;
                        }
                    }
                } finally {
                    try {
                        document.close();
                    } catch (IOException ignored) {}
                }
                if (imageNumber < iNo) {
                    System.out.printf("没有第%d张图片\n", iNo);
                }
                return null;
            }
        };
        new Thread(task).start();
    }

    @Override
    public void start(Stage stage) {
        imageFlowPane = new FlowPane();
        //File chooser
        Button imageChooser = new Button("选择PDF");
        imageChooser.setOnMouseClicked(e -> {
            FileChooser fileChooser = new FileChooser();
            fileChooser.setTitle("请选择PDF");
            fileChooser.setInitialDirectory(new File(System.getProperty("user.home") + "/Downloads")
            );
            fileChooser.getExtensionFilters().addAll(
                    new FileChooser.ExtensionFilter("PDF", "*.pdf")
            );
            File pdfFile = fileChooser.showOpenDialog(stage);
            if (null == pdfFile) {
                System.out.println("用户取消选中文件");
                return;
            }
            //清空显示在flowPane上面的图片
            Platform.runLater(() -> imageFlowPane.getChildren().clear());
            //读取一个PDF
            try {
                PDDocument document = PDDocument.load(pdfFile);
                pdfProcess(document);
            } catch (IOException ex) {
                System.out.println("加载PDF文件出错，" + ex.getMessage());
                Platform.exit();
            }
            file = pdfFile;
        });
        //页数，从0开始
        Text pageNoTip = new Text("页数（0开始）");
        TextField pageNoInput = new TextField();
        pageNo = new SimpleStringProperty();
        pageNo.bind(pageNoInput.textProperty());
        //页数，从0开始
        Text imageNoTip = new Text("图片（0开始）");
        TextField imageNoInput = new TextField();
        imageNo = new SimpleStringProperty();
        imageNo.bind(imageNoInput.textProperty());
        //触发执行
        Button runBtn = new Button("运行");
        runBtn.setOnMouseClicked(event -> {
            if (file == null) return;
            //读取一个PDF
            try {
                PDDocument document = PDDocument.load(file);
                pdfProcess(document);
            } catch (IOException ex) {
                System.out.println("加载PDF文件出错，" + ex.getMessage());
                Platform.exit();
            }
        });
        FlowPane bottomControl = new FlowPane();
        bottomControl.getChildren().addAll(imageChooser, pageNoTip, pageNoInput, imageNoTip, imageNoInput, runBtn);
        bottomControl.setPrefHeight(50);
        BorderPane root = new BorderPane();
        root.setCenter(imageFlowPane);
        root.setBottom(bottomControl);
        Scene scene = new Scene(root, DEFAULT_WIDTH, DEFAULT_HEIGHT + 50);
        stage.setScene(scene);
        stage.setTitle("PDF图片查看与切割效果检查");
        stage.setResizable(false);
        stage.show();
    }
}
